package com.study.spark.scala.ml.classification

import org.apache.spark.ml.classification.{DecisionTreeClassifier, LinearSVC, NaiveBayes}
import org.apache.spark.ml.evaluation.MulticlassClassificationEvaluator
import org.apache.spark.ml.feature.VectorAssembler
import org.apache.spark.sql.SparkSession

import scala.util.Random

/**
  * 分类算法
  *   1）朴素贝叶斯，多分类，基于概率统计
  *   2) SVM， 二分类
  *   3）决策树，多分类
  * @author stephen
  * @create 2019-04-07 15:27
  * @since 1.0.0
  */
object ClassificationDemo {
  def main(args: Array[String]): Unit = {
    val spark = SparkSession.builder()
      .appName("Classification Demo")
      .master("local[2]")
      .getOrCreate()
    val file = spark.read
      .format("csv")
      .option("sep", ",")
      .load("/Users/stephen/Documents/03code/java-demo/bigdata-study/study-spark/src/main/resource/classification/iris.data")

    val random = new Random()
    import spark.implicits._
    val data = file.map(row => {
      val label = row.getString(4) match {
        case "Iris-setosa" => 0
        case "Iris-versicolor" => 1
        case "Iris-virginica" => 2
      }
      (
        row.getString(0).toDouble,
        row.getString(1).toDouble,
        row.getString(2).toDouble,
        row.getString(3).toDouble,
        label,
        random.nextDouble()
      )
    }).toDF("_c0", "_c1", "_c2", "_c3", "label", "random").sort("random")//.where("label=0 or label=1")

    val assembler = new VectorAssembler()
      .setInputCols(Array("_c0", "_c1", "_c2", "_c3"))
      .setOutputCol("features")
    val dataset = assembler.transform(data)
    val Array(train,test) = dataset.randomSplit(Array(0.8,0.2))
    // 朴素贝叶斯
    /**
    val bayes = new NaiveBayes().setFeaturesCol("features").setLabelCol("label")
    val bayesModel = bayes.fit(train)
    bayesModel.transform(test).show()
    **/

    // SVM
    /**
    val svm = new LinearSVC().setFeaturesCol("features").setLabelCol("label")
      .setMaxIter(20).setRegParam(0.1)
    val svmModel = svm.fit(train)
    svmModel.transform(test).show()
    **/

    val dt = new DecisionTreeClassifier().setFeaturesCol("features").setLabelCol("label")
    val dtModel = dt.fit(train)
    val result = dtModel.transform(test)
    result.show()
    // 评估器
    val evaluator = new MulticlassClassificationEvaluator()
      .setLabelCol("label")
      .setPredictionCol("prediction")
      .setMetricName("accuracy")
    val accuracy = evaluator.evaluate(result)
    println(s"accuracy is $accuracy")
  }
}
